# -*- coding: utf-8 -*-
'''
Created on Aug 27, 2019

@author: yl
'''
from scipy import signal
from scipy.optimize import curve_fit
from scipy.signal import *
from FFT_Interpolation import *
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def fit_func(x, a, b, c):
    return a*(x-b)**2 + c

def w_z(w_0, Z_R, Z):
    return w_0*np.sqrt(1 + (Z/Z_R)**2)

def R_z(Z_R, Z):
    return Z*(1+(Z_R/Z)**2)

def Phi_Gouy(Z_R, Z):
    return np.arctan(Z/Z_R)    

img_set = []
hor_phase_centers = []
hor_freq_set = []
hor_f_fit_set = []
hor_phi_fit_set = []
angle_hor_set = []

j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]
k = 2 * np.pi / Lamda
fs_cam = 50 ### Hardware Trigger

pix_size = 5.3e-6
pix_num = 1280
screen_diameter = pix_num * pix_size

window = signal.gaussian(pix_num, std=pix_num/10)
window_y = signal.gaussian(1024, std=pix_num/10)
window_x, window_y = np.meshgrid(window, window_y)
window_2d = window_x*window_y

I_0 = 127
V_r_x, V_r_y, V_r_z = 0.000, 0.0000, 1 # Reference
V_m_x, V_m_y, V_z_m = 0.001, 0.001, 1 # Measurement
L = 0.1
M, N = 0.1, 0.05
w_0 = 1e-3
Z_R = np.pi * w_0**2 / Lamda
dx = np.linspace(0, (pix_num-1)*pix_size, num=pix_num)
dy = np.linspace(0, (1024-1)*pix_size, num=1024)
X, Y = np.meshgrid(dx, dy)

frame_num = 10
line_num = 2
D_set = np.arange(10)*Lamda/5

hor_f_lines = [[None]*frame_num]*line_num
hor_phi_lines = [[None]*frame_num]*line_num
# print(hor_f_lines)

for frame in range(frame_num):
    print(frame)
    D = D_set[frame]
    
    Z_p_r = M + N + L * (2/(1-V_r_x**2-V_r_y**2) - 1)
    Z_p_m = M + N + (L+D) * (2/(1-V_m_x**2-V_m_y**2) - 1)
    
    diff_Z_p = D * 2 / (1-V_m_x**2-V_m_y**2) + L * 2 * ((V_m_x**2+V_m_y**2)-(V_r_x**2+V_r_y**2)) / (1-V_m_x**2-V_m_y**2) / (1-V_r_x**2-V_r_y**2)
    
    d_mea = (X*2*V_m_x+Y*2*V_m_y)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_ref = (X*2*V_r_x+Y*2*V_r_y)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
    
    r_mea = np.sqrt( X**2 + Y**2 + (L+D)**2 - ((2*X*V_m_x+2*Y*V_m_y+(L+D)*(1-V_m_x**2-V_m_y**2))/(1+V_m_x**2+V_m_y**2))**2 )
    r_ref = np.sqrt( X**2 + Y**2 + (L)**2 - ((2*X*V_r_x+2*Y*V_r_y+L*(1-V_r_x**2-V_r_y**2))/(1+V_r_x**2+V_r_y**2))**2 )
    
    R_ref = R_z(Z_R, Z_p_r+d_ref)
    R_mea = R_z(Z_R, Z_p_m+d_mea)
    
    Phi_gouy_ref = Phi_Gouy(Z_R, Z_p_r+d_ref)
    Phi_gouy_mea = Phi_Gouy(Z_R, Z_p_m+d_mea)
    
    diff_phi = k * (diff_Z_p + (d_mea-d_ref) + r_mea**2/2/R_mea - r_ref**2/2/R_ref) + Phi_gouy_mea - Phi_gouy_ref
    A_beat = 0.5 * I_0 * w_0**2 / w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea) * np.exp((-X**2-Y**2)/w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea))
#     A_beat = 0.5 * I_0
    I_beat = A_beat*(1+np.cos(diff_phi))
    
    R = 2
    center = [(pix_num-1-800)*pix_size/2, (pix_num+800)*pix_size/2]
    r = np.sqrt((X-center[0])**2 + (Y-center[1])**2)
    d = R - np.sqrt(R**2 - r**2)
    A = 2 * np.pi * (d+2*D) / Lamda
    I_r_2 = np.sin(A)**2 * I_0/2
    
    I_beat = I_beat + I_r_2
    I_beat = I_beat * window_2d
    I_beat = I_beat.astype(np.int)
#     I_beat = I_beat
    for line in range(line_num):
        hor_center = np.array(I_beat[int(512+(line)*33)])
        ver_center = I_beat[:,int(640+(line-1)*100)]
        DC_num = 1000
        hor_freq_estim, hor_phase_estim, hor_freqline, hor_sig_magnitude, hor_sig_phase,  hor_m_k_num, hor_X_m_k, hor_freq_for_phase = FFT_interpolation_2(hor_center, pix_size, 1e5, DC_num)
        ver_freq_estim, ver_phase_estim, ver_freqline, ver_sig_magnitude, ver_sig_phase, ver_m_k_num, ver_X_m_k, ver_freq_for_phase = FFT_interpolation_2(ver_center, pix_size, 1e5, DC_num)
        
        hor_FFT_start = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][0]+DC_num
        hor_FFT_end = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][-1]+DC_num
        hor_fit_x = hor_freqline[hor_FFT_start:hor_FFT_end+1]
        hor_fit_y = hor_sig_magnitude[hor_FFT_start:hor_FFT_end+1]
        hor_fit_phase = hor_sig_phase[hor_FFT_start:hor_FFT_end+1]
        hor_params = curve_fit(fit_func, hor_fit_x, hor_fit_y)
        [hor_a, hor_b, hor_c] = hor_params[0]
        hor_f_fit = hor_b
        
        hor_f_lines[line][frame] = hor_f_fit
        hor_phi_lines[line][frame] = hor_phase_estim
        
        ver_FFT_start = np.where(ver_sig_magnitude[DC_num:] > ver_X_m_k*0.4)[0][0]+DC_num
        ver_FFT_end = np.where(ver_sig_magnitude[DC_num:] > ver_X_m_k*0.4)[0][-1]+DC_num
        ver_fit_x = ver_freqline[ver_FFT_start:ver_FFT_end+1]
        ver_fit_freq = ver_sig_magnitude[ver_FFT_start:ver_FFT_end+1]
        ver_params = curve_fit(fit_func, ver_fit_x, ver_fit_freq)
        [ver_a, ver_b, ver_c] = ver_params[0]
        ver_f_fit = ver_b
        
#         ver_freqs.append(ver_f_fit)
#         ver_phases.append(ver_phase_estim)
#     print(hor_f_lines)  

plt.figure(1)
#     im = plt.imshow(I_beat, cmap='gray')
#     plt.colorbar(im, fraction=0.046, pad=0.04)
plt.plot(hor_f_lines[0], 'b')
plt.plot(hor_f_lines[1], 'r')
plt.grid(which='both', axis='both')
plt.show()
